#!/usr/bin/env python3
"""
从 extracted_queries.txt 读取查询并执行
参考 model_learning.py 的执行方式
"""

from TLSMapper.TLSSUT import *
from TLSMapper.TLSProtocol import *
import ast
import argparse
import sys

# 密码套件配置（与 model_learning.py 保持一致）
ciphersuites = [
    CipherSuite.TLS_AES_128_GCM_SHA256,
    CipherSuite.TLS_AES_256_GCM_SHA384,
    CipherSuite.TLS_ECDHE_ECDSA_WITH_AES_256_GCM_SHA384,
    CipherSuite.TLS_DHE_RSA_WITH_CHACHA20_POLY1305_SHA256,
    CipherSuite.TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384
]

# 证书路径配置（与 model_learning_tlse.py 保持一致）
CERT_BASE = '/home/ubuntu/experiments/certs'
CLIENT_KEY = f'{CERT_BASE}/client-key.pem'
CLIENT_CERT = f'{CERT_BASE}/client-cert.pem'
SERVER_CERT = f'{CERT_BASE}/server-cert.pem'
CA_CERT = f'{CERT_BASE}/ca-cert.pem'
SERVER_HOST = '127.0.0.1'
SERVER_PORT = 1111


def parse_query(query_str):
    """
    解析查询字符串（可能是元组或列表格式）
    
    Args:
        query_str: 查询字符串，如 "('ClosureAlert', 'ClientHello')" 或 "['ClientHello', 'Finish']"
    
    Returns:
        列表格式的查询
    """
    try:
        # 使用 ast.literal_eval 安全地解析字符串
        parsed = ast.literal_eval(query_str)
        # 转换为列表（如果是元组）
        if isinstance(parsed, tuple):
            return list(parsed)
        elif isinstance(parsed, list):
            return parsed
        else:
            raise ValueError(f"Unexpected query format: {query_str}")
    except Exception as e:
        print(f"错误: 无法解析查询 '{query_str}': {e}")
        return None


def execute_query(sul, query, query_num):
    """
    执行单个查询
    
    Args:
        sul: TLSSUT 对象
        query: 查询列表
        query_num: 查询编号
    
    Returns:
        (success, output) 元组
    """
    try:
        print(f"\n[{query_num}] 执行查询: {query}")
        print("-" * 80)
        
        # 执行查询
        output = sul.query(query)
        
        print(f"响应: {output}")
        print("=" * 80)
        
        return (True, output)
    except Exception as e:
        print(f"✗ 查询执行失败: {e}")
        import traceback
        traceback.print_exc()
        print("=" * 80)
        return (False, None)


def main():
    parser = argparse.ArgumentParser(
        description='从 extracted_queries.txt 读取查询并执行',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 执行所有查询
  python3 execute_queries.py --input extracted_queries.txt
  
  # 从第10行执行到第20行（包含）
  python3 execute_queries.py --input extracted_queries.txt --start 10 --end 20
  
  # 从第1行执行到第100行
  python3 execute_queries.py --input extracted_queries.txt --start 1 --end 100
  
  # 只执行第50行
  python3 execute_queries.py --input extracted_queries.txt --start 50 --end 50
        """
    )
    
    parser.add_argument(
        '--input',
        type=str,
        default='extracted_queries.txt',
        help='输入文件路径（默认: extracted_queries.txt）'
    )
    
    parser.add_argument(
        '--start',
        type=int,
        default=1,
        help='起始行号（从1开始，默认: 1）'
    )
    
    parser.add_argument(
        '--end',
        type=int,
        default=None,
        help='结束行号（包含，默认: 文件末尾）'
    )
    
    parser.add_argument(
        '--implementation',
        type=str,
        default='tlse',
        choices=['tlse', 'wolfssl', 'openssl', 'mbedtls', 'gnutls', 'matrixssl', 'boringssl'],
        help='TLS 实现（默认: tlse）'
    )
    
    parser.add_argument(
        '--version',
        type=str,
        default='1.2',
        choices=['1.2', '1.3'],
        help='TLS 版本（默认: 1.2）'
    )
    
    parser.add_argument(
        '--mutual',
        action='store_true',
        help='使用双向认证（需要客户端证书）'
    )
    
    args = parser.parse_args()
    
    # 读取查询文件
    try:
        with open(args.input, 'r', encoding='utf-8') as f:
            lines = f.readlines()
    except FileNotFoundError:
        print(f"错误: 文件 '{args.input}' 不存在")
        return
    except Exception as e:
        print(f"错误: 无法读取文件 '{args.input}': {e}")
        return
    
    # 确定执行范围
    start_line = max(1, args.start) - 1  # 转换为0-based索引
    end_line = args.end if args.end is not None else len(lines)
    end_line = min(end_line, len(lines))  # 确保不超过文件长度
    
    if start_line >= len(lines):
        print(f"错误: 起始行号 {args.start} 超出文件范围（文件共 {len(lines)} 行）")
        return
    
    if start_line >= end_line:
        print(f"错误: 起始行号 {args.start} 必须小于结束行号 {end_line}")
        return
    
    print("=" * 80)
    print(f"读取查询文件: {args.input}")
    print(f"执行范围: 第 {args.start} 行到第 {end_line} 行（共 {end_line - start_line} 个查询）")
    print(f"TLS 实现: {args.implementation}")
    print(f"TLS 版本: {args.version}")
    print(f"认证模式: {'双向认证' if args.mutual else '单向认证'}")
    print("=" * 80)
    
    # 解析 TLS 版本
    if args.version == '1.2':
        tls_version = (3, 3)
    else:
        tls_version = (3, 4)
    
    # 创建 TLS 协议配置
    tlsimp = TLSProtocol(
        version=tls_version,
        cipher_suite=ciphersuites,
        implementation=args.implementation,
        verify=False,
        fuzz_tls=(False, None, None),
        target=(SERVER_HOST, SERVER_PORT)
    )
    
    # 创建 SUT
    if args.mutual:
        sul = TLSSUT(
            TLSpro=tlsimp,
            keyfile=CLIENT_KEY,
            certfile=CLIENT_CERT,
            server_cert=CA_CERT
        )
    else:
        sul = TLSSUT(
            TLSpro=tlsimp,
            keyfile=None,
            certfile=None,
            server_cert=CA_CERT
        )
    
    # 执行查询
    results = []
    success_count = 0
    fail_count = 0
    
    for i in range(start_line, end_line):
        line = lines[i].strip()
        if not line:  # 跳过空行
            continue
        
        query_num = i + 1  # 行号（从1开始）
        
        # 解析查询
        query = parse_query(line)
        if query is None:
            fail_count += 1
            results.append((query_num, line, False, None))
            continue
        
        # 执行查询
        success, output = execute_query(sul, query, query_num)
        
        if success:
            success_count += 1
        else:
            fail_count += 1
        
        results.append((query_num, query, success, output))
    
    # 打印总结
    print("\n" + "=" * 80)
    print("执行总结")
    print("=" * 80)
    print(f"总查询数: {len(results)}")
    print(f"成功: {success_count}")
    print(f"失败: {fail_count}")
    print("=" * 80)
    
    # 保存结果到文件
    output_file = f"query_results_{args.start}_{end_line}.txt"
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(f"查询执行结果（第 {args.start} 行到第 {end_line} 行）\n")
        f.write("=" * 80 + "\n")
        for query_num, query, success, output in results:
            status = "✓ 成功" if success else "✗ 失败"
            f.write(f"\n[{query_num}] {status}\n")
            f.write(f"查询: {query}\n")
            if output:
                f.write(f"响应: {output}\n")
            f.write("-" * 80 + "\n")
    
    print(f"\n结果已保存到: {output_file}")


if __name__ == '__main__':
    main()

